import numpy as np
import cv2
from sklearn.metrics import confusion_matrix

def calculate_iou(pred_mask, true_mask):
    """
    计算IoU (Intersection over Union)
    
    Args:
        pred_mask (numpy.ndarray): 预测的二值掩码
        true_mask (numpy.ndarray): 真实的二值掩码
    
    Returns:
        float: IoU值
    """
    # 确保掩码是二值的
    pred_mask = (pred_mask > 0).astype(np.uint8)
    true_mask = (true_mask > 0).astype(np.uint8)
    
    # 计算交集和并集
    intersection = np.logical_and(pred_mask, true_mask).sum()
    union = np.logical_or(pred_mask, true_mask).sum()
    
    # 避免除零错误
    if union == 0:
        return 0.0
    
    iou = intersection / union
    return float(iou)

def calculate_f1_score(pred_mask, true_mask):
    """
    计算F1分数
    
    Args:
        pred_mask (numpy.ndarray): 预测的二值掩码
        true_mask (numpy.ndarray): 真实的二值掩码
    
    Returns:
        float: F1分数
    """
    # 确保掩码是二值的
    pred_mask = (pred_mask > 0).astype(np.uint8)
    true_mask = (true_mask > 0).astype(np.uint8)
    
    # 计算混淆矩阵
    tn, fp, fn, tp = confusion_matrix(true_mask.flatten(), pred_mask.flatten(), labels=[0, 1]).ravel()
    
    # 计算精确率和召回率
    precision = tp / (tp + fp) if (tp + fp) > 0 else 0
    recall = tp / (tp + fn) if (tp + fn) > 0 else 0
    
    # 计算F1分数
    f1 = 2 * (precision * recall) / (precision + recall) if (precision + recall) > 0 else 0
    
    return float(f1)

def calculate_metrics(pred_mask_path, true_mask_path):
    """
    计算预测掩码相对于真实掩码的评估指标
    
    Args:
        pred_mask_path (str): 预测掩码图像的路径
        true_mask_path (str): 真实掩码图像的路径
    
    Returns:
        dict: 包含各种评估指标的字典
    """
    # 读取掩码图像
    pred_mask = cv2.imread(pred_mask_path, cv2.IMREAD_GRAYSCALE)
    true_mask = cv2.imread(true_mask_path, cv2.IMREAD_GRAYSCALE)
    
    # 确保两个掩码尺寸相同
    if pred_mask.shape != true_mask.shape:
        true_mask = cv2.resize(true_mask, (pred_mask.shape[1], pred_mask.shape[0]), interpolation=cv2.INTER_NEAREST)
    
    # 计算IoU
    iou = calculate_iou(pred_mask, true_mask)
    
    # 计算F1分数
    f1 = calculate_f1_score(pred_mask, true_mask)
    
    # 计算准确率
    total_pixels = pred_mask.size
    correct_pixels = np.sum(pred_mask == true_mask)
    accuracy = correct_pixels / total_pixels
    
    # 返回所有指标
    metrics = {
        'iou': iou,
        'f1': f1,
        'accuracy': accuracy
    }
    
    return metrics

def compare_models(model_results):
    """
    比较不同模型的性能
    
    Args:
        model_results (dict): 包含不同模型评估结果的字典
                             格式: {'model_name': {'iou': value, 'f1': value, ...}}
    
    Returns:
        dict: 包含模型比较结果的字典
    """
    # 计算平均指标
    avg_metrics = {}
    for model_name, metrics in model_results.items():
        avg_iou = np.mean([result['iou'] for result in metrics])
        avg_f1 = np.mean([result['f1'] for result in metrics])
        
        avg_metrics[model_name] = {
            'avg_iou': avg_iou,
            'avg_f1': avg_f1
        }
    
    # 找出最佳模型
    best_iou_model = max(avg_metrics, key=lambda x: avg_metrics[x]['avg_iou'])
    best_f1_model = max(avg_metrics, key=lambda x: avg_metrics[x]['avg_f1'])
    
    comparison = {
        'avg_metrics': avg_metrics,
        'best_iou_model': best_iou_model,
        'best_f1_model': best_f1_model
    }
    
    return comparison 